#include "UI.h"

map<Core*, UI*> UI::instances;

UI::UI(HINSTANCE hInstance)
{
	unauthorizedAccessCount = 0;

	this->hInstance = hInstance;

	WM_TASKBARCREATED = RegisterWindowMessage(TEXT("TaskbarCreated"));
	
	Window::Create(hInstance, TEXT("AuroraWindow"), WndProc, 0, WS_POPUP, CS_NOCLOSE);
	
	CreateTrayMenu();
	CheckMenuItem(trayMenu, TrayMenu_StartWithWindows, IsAutorunEnabled() ? MF_CHECKED : MF_UNCHECKED);

	auroraCore = new Core(hInstance);
	instances[auroraCore] = this;

	auroraCore->ActiveChanged(CoreActiveChanged);
	auroraCore->LockedChanged(CoreLockedChanged);
	auroraCore->LockAtUnplugChanged(CoreLockAtUnplugChanged);
	auroraCore->KeyDevicesChanged(CoreKeyDevicesChanged);
	auroraCore->UnauthorizedAccess(CoreUnauthorizedAccess);

	registerDriveWindow = new RegisterDriveWindow(this);
	registerDriveWindow->Create(hInstance, LoadIcon(hInstance, MAKEINTRESOURCE(IDI_ICON1)));
	registerDriveWindow->DrivePluggedIn(DrivePluggedIn);

	Settings::Load();

	auroraCore->SetKeyDeviceCollection(Settings::GetData()->KeyDevices);
	auroraCore->SetLockAtUnplug(Settings::GetData()->LockAtUnplug);

	if (Settings::GetData()->Active)
		auroraCore->Activate(Settings::GetData()->Locked);
}

UI::~UI()
{
	DestroyMenu(trayMenu);
	Shell_NotifyIcon(NIM_DELETE, &notifyIconData);

	instances.erase(auroraCore);

	delete registerDriveWindow;
	registerDriveWindow = 0;

	delete auroraCore;
	auroraCore = 0;
}


void UI::CreateTrayMenu()
{
	notifyIconData.cbSize = sizeof(NOTIFYICONDATA);
	notifyIconData.hIcon = LoadIcon(hInstance, MAKEINTRESOURCE(IDI_ICON1));
	notifyIconData.uID = 1;
	notifyIconData.uVersion = NOTIFYICON_VERSION_4;
	notifyIconData.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
	_tcscpy_s(notifyIconData.szTip, TEXT("Aurora"));
	notifyIconData.hWnd = this->GetHWnd();
	notifyIconData.uCallbackMessage = WM_APP + 1;
	Shell_NotifyIcon(NIM_ADD, &notifyIconData);

	trayMenu = CreatePopupMenu();
	
	AppendMenu(trayMenu, MF_STRING | MF_DISABLED, TrayMenu_Active, TEXT("Active"));
	AppendMenu(trayMenu, MF_SEPARATOR, 0, 0);
	AppendMenu(trayMenu, MF_STRING, TrayMenu_RegisterKeyDevice, TEXT("Register KeyDevice"));
	
	AppendMenu(trayMenu, MF_STRING | MF_DISABLED, 0, TEXT("KeyDevices"));
	
	AppendMenu(trayMenu, MF_SEPARATOR, 0, 0);
	AppendMenu(trayMenu, MF_STRING, TrayMenu_LockAtUnplug, TEXT("Lock at unplug"));
	AppendMenu(trayMenu, MF_STRING, TrayMenu_StartWithWindows, TEXT("Start with Windows"));
	AppendMenu(trayMenu, MF_SEPARATOR, 0, 0);
	AppendMenu(trayMenu, MF_STRING, TrayMenu_Exit, TEXT("Exit"));

	SetMenuDefaultItem(trayMenu, TrayMenu_Active, 0);
}

void UI::ToggleCoreActiveState()
{
	if (auroraCore->IsActive())
		auroraCore->Deactivate();
	else
		auroraCore->Activate();
}

bool UI::IsAutorunEnabled()
{
	HKEY hKey = 0;
	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, TEXT("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run"), 0, KEY_READ | KEY_QUERY_VALUE, &hKey) == ERROR_SUCCESS)
	{
		LSTATUS result = RegQueryValueEx(hKey, TEXT("Aurora"), 0, 0, 0, 0);
		
		if (hKey)
			RegCloseKey(hKey);

		if (result == ERROR_SUCCESS)
			return true;
		else
			return false;
	}
	else
	{
		return false;
	}
}

bool UI::EnableAutorun(bool value)
{
	TCHAR enableAutorun[20] = {0};

	if (value)
		_stprintf_s(enableAutorun, TEXT("%s"), TEXT("true"));
	else
		_stprintf_s(enableAutorun, TEXT("%s"), TEXT("false"));

	TCHAR path[MAX_PATH];
	GetModuleFileNameW(GetModuleHandle(0), path, MAX_PATH);

	TCHAR args[512] = {0};
	_stprintf_s(args, TEXT("%s Aurora \"%s\""), enableAutorun, path);
	int result = (int)ShellExecute(this->GetHWnd(), 0, TEXT("SetGlobalAutorun.exe"), args, 0, SW_HIDE);

	return result > 32;		
}

void UI::WriteLog()
{
	SYSTEMTIME time;
	GetLocalTime(&time);
	ofstream file("UnauthorizedAccess.log", ios_base::app);
	file << time.wDay << "." << time.wMonth << "." << time.wYear << " -- " << time.wHour << ":" << time.wMinute << ":" << time.wSecond << endl; 
	file.close();
}


//is this necessary instead of doing that in the constructor?
//void UI::Create(HINSTANCE hInstance)
//{
//Window::Create(hInstance, TEXT("AuroraWindow"), WndProc, 0, WS_POPUP, CS_NOCLOSE);
//}


void UI::CoreActiveChanged(Core* sender, bool value)
{
	if (value)
		instances[sender]->notifyIconData.hIcon = LoadIcon(instances[sender]->hInstance, MAKEINTRESOURCE(IDI_ICON2));
	else
		instances[sender]->notifyIconData.hIcon = LoadIcon(instances[sender]->hInstance, MAKEINTRESOURCE(IDI_ICON1));

	Shell_NotifyIcon(NIM_MODIFY, &instances[sender]->notifyIconData);

	CheckMenuItem(instances[sender]->trayMenu, TrayMenu_Active, value ? MF_CHECKED : MF_UNCHECKED);

	Settings::GetData()->Active = value;	
	Settings::Save();
}

void UI::CoreLockAtUnplugChanged(Core* sender, bool value)
{
	CheckMenuItem(instances[sender]->trayMenu, TrayMenu_LockAtUnplug, value ? MF_CHECKED : MF_UNCHECKED);

	Settings::GetData()->LockAtUnplug = value;
	Settings::Save();
}

void UI::CoreLockedChanged(Core* sender, bool value)
{
	if (!value)
	{
		if (instances[sender]->unauthorizedAccessCount > 0)
		{
			PlaySound(0, 0, 0);

			TCHAR message[100] = {0};
			_stprintf_s(message, TEXT("Unauthorized Access: %d"), instances[sender]->unauthorizedAccessCount);
			MessageBox(instances[sender]->GetHWnd(), message, TEXT("Unauthorized Access - Aurora"), MB_OK | MB_ICONINFORMATION | MB_DEFBUTTON1 | MB_TOPMOST);
		}

		instances[sender]->unauthorizedAccessCount = 0;
	}

	Settings::GetData()->Locked = value;
	Settings::Save();
}

void UI::RefreshKeyDevicesMenu()
{
	DestroyMenu(keyDevicesSubMenu);
	keyDevicesSubMenu = CreatePopupMenu();
	
	int counter = 0;

	map<DWORD, tstring> keyDevices = auroraCore->GetKeyDeviceCollection();
	map<DWORD, tstring>::iterator keyDevice;
	for (keyDevice = keyDevices.begin(); keyDevice != keyDevices.end(); keyDevice++)
	{
		AppendMenu(keyDevicesSubMenu, MF_STRING, TrayMenu_KeyDevicesStart + counter, keyDevice->second.c_str());
		
		MENUITEMINFO mii = {0};
		mii.cbSize = sizeof(MENUITEMINFO);
		mii.fMask = MIIM_DATA;
		mii.dwItemData = keyDevice->first;
		SetMenuItemInfo(keyDevicesSubMenu, TrayMenu_KeyDevicesStart + counter, FALSE, &mii);

		counter++;
	}
	AppendMenu(keyDevicesSubMenu, MF_SEPARATOR, 0, 0);
	AppendMenu(keyDevicesSubMenu, MF_STRING, TrayMenu_ClearAllKeyDevices, TEXT("Clear all KeyDevices"));

	ModifyMenu(trayMenu, 3, MF_BYPOSITION | MF_POPUP | MF_STRING, (UINT_PTR)keyDevicesSubMenu, TEXT("KeyDevices"));
}

void UI::CoreKeyDevicesChanged(Core* sender, size_t value)
{
	if (value > 0) //if there is no key device, the menu is hidden so no refresh is necessary
		instances[sender]->RefreshKeyDevicesMenu();

	EnableMenuItem(instances[sender]->trayMenu, 3, MF_BYPOSITION | (value > 0 ? MF_ENABLED : MF_DISABLED));
	EnableMenuItem(instances[sender]->trayMenu, TrayMenu_Active, value > 0 ? MF_ENABLED : MF_DISABLED);

	Settings::GetData()->KeyDevices = sender->GetKeyDeviceCollection();
	Settings::Save();
}

void UI::CoreUnauthorizedAccess(Core* sender)
{
	//option: take a photo, if "TakePicture.exe" exists. If not, ShellExecute will return ERROR_FILE_NOT_FOUND.
	ShellExecute(instances[sender]->GetHWnd(), 0, TEXT("TakePicture.exe"), 0, 0, SW_HIDE);

	//option: send an e-mail
	//ShellExecute(instances[sender]->GetHWnd(), 0, TEXT("<filename>"), 0, 0, SW_HIDE);

	//option: play an alarm sound
	if (Utilities::ExistsFile(TEXT("alert.wav")))
		PlaySound(TEXT("alert.wav"), instances[sender]->hInstance, SND_FILENAME | SND_ASYNC);


	instances[sender]->unauthorizedAccessCount++;
	instances[sender]->WriteLog();
}


void UI::DrivePluggedIn(RegisterDriveWindow* sender, TCHAR driveLetter, DWORD driveSerial, LPVOID parent, LPVOID tag)
{
	int result = ((UI*)parent)->auroraCore->AddKeyDevice(driveSerial, Utilities::GetVolumeName(driveLetter));

	if (result == KEYDEVICE_INVALIDSERIAL)
		MessageBox(((UI*)parent)->GetHWnd(), TEXT("Invalid device."), TEXT("Error - Aurora"), MB_OK | MB_ICONERROR | MB_DEFBUTTON1 | MB_TOPMOST);
	else if (result == KEYDEVICE_ALREADYEXISTS)
		MessageBox(((UI*)parent)->GetHWnd(), TEXT("The specified device is already registered."), TEXT("Error - Aurora"), MB_OK | MB_ICONERROR | MB_DEFBUTTON1 | MB_TOPMOST);
}

LRESULT CALLBACK UI::WndProc(LPVOID sender, HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{	
	UI* instance = (UI*)sender;

	if (uMsg == instance->WM_TASKBARCREATED)
		Shell_NotifyIcon(NIM_ADD, &instance->notifyIconData);

	switch (uMsg)
	{
	case WM_APP + 1:
		{
			switch(LOWORD(lParam))
			{
			case WM_RBUTTONDOWN:
			case WM_CONTEXTMENU:
				{
					POINT cursorPos = Utilities::GetCursorPosition();
					SetForegroundWindow(hWnd);
					TrackPopupMenu(instance->trayMenu, TPM_VERPOSANIMATION | TPM_LEFTBUTTON, cursorPos.x, cursorPos.y, 0, hWnd, 0);
					PostMessage(hWnd, WM_NULL, 0, 0);
					break;
				}
			case WM_LBUTTONUP:
				{
					instance->ToggleCoreActiveState();
					break;
				}
			}

			break;
		}
	case WM_COMMAND:
		{
			WORD menuID = LOWORD(wParam);
			if (menuID >= TrayMenu_KeyDevicesStart)
			{
				MENUITEMINFO mii = {0};
				mii.cbSize = sizeof(MENUITEMINFO);
				mii.fMask = MIIM_DATA;
				GetMenuItemInfo(instance->keyDevicesSubMenu, menuID, FALSE, &mii);

				TCHAR menuString[_MAX_PATH + 1] = {0};
				GetMenuString(instance->keyDevicesSubMenu, menuID, menuString, _MAX_PATH, 0);

				TCHAR message[512] = {0};
				_stprintf_s(message, TEXT("Are you sure you want to delete the KeyDevice \"%s\"?"), menuString);

				if (MessageBox(hWnd, message, TEXT("Delete KeyDevice - Aurora"), MB_YESNO | MB_ICONQUESTION | MB_DEFBUTTON2 | MB_TOPMOST) == IDYES)
				{
					instance->auroraCore->RemoveKeyDevice((DWORD)mii.dwItemData);
				}
			}
			else
			{
				switch (menuID)
				{
				case TrayMenu_Active:
					{
						instance->ToggleCoreActiveState();
						break;
					}
				case TrayMenu_RegisterKeyDevice:
					{
						instance->registerDriveWindow->Show();
						break;
					}
				case TrayMenu_ClearAllKeyDevices:
					{
						if (MessageBox(hWnd, TEXT("Are you sure?"), TEXT("Clear all KeyDevices - Aurora"), MB_YESNO | MB_ICONQUESTION | MB_DEFBUTTON2 | MB_TOPMOST) == IDYES)
							instance->auroraCore->ClearKeyDevices();

						break;
					}
				case TrayMenu_LockAtUnplug:
					{
						instance->auroraCore->SetLockAtUnplug(!instance->auroraCore->GetLockAtUnplug());
						CheckMenuItem(instance->trayMenu, TrayMenu_LockAtUnplug, instance->auroraCore->GetLockAtUnplug() ? MF_CHECKED : MF_UNCHECKED);
						break;
					}
				case TrayMenu_StartWithWindows:
					{
						bool isAutorunEnabled = instance->IsAutorunEnabled();

						if (!instance->EnableAutorun(!isAutorunEnabled))
							MessageBox(hWnd, TEXT("Error"), TEXT("Error - Aurora"), MB_OK | MB_ICONERROR | MB_DEFBUTTON1 | MB_TOPMOST);
						else
							CheckMenuItem(instance->trayMenu, TrayMenu_StartWithWindows, !isAutorunEnabled ? MF_CHECKED : MF_UNCHECKED);

						break;
					}
				case TrayMenu_Exit:
					{
						PostQuitMessage(0);
						break;
					}			
				}
			}

			break;
		}
	case WM_DESTROY:	
		{
			PostQuitMessage(0);
			break;
		}
	}

	return DefWindowProc(hWnd, uMsg, wParam, lParam);
}